from QKDNetwork import QKDNetwork
from docplex.mp.model import Model
import random


def provision_with_max_key_num(routeList: list, net: QKDNetwork):
    # 每次供应最大数量的密钥
    model = Model(name="provision_with_max_key_num")

    def get_sd_idx(sd_start, sd_end):
        for idx, i in enumerate(net.sd_list):
            if i[0] == sd_start and i[1] == sd_end:
                return idx
        return -1

    D = model.continuous_var(name="D", lb=0)

    # 满足了多少密钥
    x_dim = len(routeList) + 1
    x = model.continuous_var_dict([i for i in range(x_dim)], lb=0, name="x")

    e_dim1, e_dim2 = net.num_nodes + 1, net.num_nodes + 1

    # 是否 u 向 v 发密钥
    e = model.binary_var_dict(
        [(u, v) for u in range(e_dim1) for v in range(e_dim2)], name="e"
    )

    def is_path_include_link(path, u, v):
        for i in range(len(path) - 1):
            if path[i] == u and path[i + 1] == v:
                return 1
            if path[i] == v and path[i + 1] == u:
                return 1
        return 0

    def u_v_in_all_paths(u, v):
        for i in routeList:
            if is_path_include_link(i["path"], u, v) == 1:
                return True
        return False

    # 选一个发射方向
    for u in net.G.nodes:
        if net.G.nodes[u]["transmitter"]:
            sum_e = model.linear_expr()
            for v in net.G.neighbors(u):
                sum_e += e[u, v]
            model.add_constraint(sum_e == 1)
        else:
            sum_e = model.linear_expr()
            for v in net.G.neighbors(u):
                sum_e += e[u, v]
            model.add_constraint(sum_e == 0)

    for idx, i in enumerate(routeList):
        remain_request_key = i["sd"][2] - i["sd"][3]
        if remain_request_key >= 0:
            model.add_constraint(x[idx] <= i["sd"][2] - i["sd"][3])
        else:
            model.add_constraint(x[idx] == 0)

    # 压链路
    for edge in net.G.edges:
        u, v = edge
        sum_T = model.linear_expr()

        for idx, i in enumerate(routeList):
            sum_T += is_path_include_link(i["path"], u, v) * x[idx]

        sum_T -= net.G[u][v]["remain_key_num"]
        if net.G.nodes[u]["transmitter"] == 1:
            sum_T -= e[u, v] * net.G.nodes[u]["transmitter_rate"]
        if net.G.nodes[v]["transmitter"] == 1:
            sum_T -= e[v, u] * net.G.nodes[v]["transmitter_rate"]

        model.add_constraint(sum_T <= 0)

    total_key_num = model.linear_expr()
    for idx, i in enumerate(routeList):
        total_key_num += (len(i["path"]) - 1) * x[idx]

    model.add_constraint(total_key_num >= D)

    model.maximize(D)

    model.export_as_lp(path="provision_with_max_key_num.lp")

    sol = model.solve()

    ret_arr = []

    if sol:

        for edge in net.G.edges:
            u, v = edge

            if (u, v) in e and e[u, v].solution_value >= 0.99:
                ret_arr.append(([u, v], net.G.nodes[u]["transmitter_rate"]))
                net.G[u][v]["remain_key_num"] += net.G.nodes[u]["transmitter_rate"]

            if (v, u) in e and e[v, u].solution_value >= 0.99:
                ret_arr.append(([v, u], net.G.nodes[v]["transmitter_rate"]))
                net.G[v][u]["remain_key_num"] += net.G.nodes[v]["transmitter_rate"]

        for i in range(len(routeList)):
            path = routeList[i]["path"]
            consume_key_num = x[i].solution_value

            if consume_key_num == 0:
                continue

            for j in range(len(path) - 1):
                link_start, link_end = path[j], path[j + 1]
                net.G[link_start][link_end]["remain_key_num"] -= consume_key_num

            sd_start, sd_end = routeList[i]["sd"][0], routeList[i]["sd"][1]
            net_sd_idx = get_sd_idx(sd_start, sd_end)

            old_sd_item = net.sd_list[net_sd_idx]

            tmp_sd_item = (
                old_sd_item[0],
                old_sd_item[1],
                old_sd_item[2],
                old_sd_item[3] + consume_key_num,
            )

            net.sd_list[net_sd_idx] = tmp_sd_item
            routeList[i]["sd"] = tmp_sd_item

        # print("ret_arr")
        # for i in ret_arr:
        #     print(i)


def provision_with_noise_lp(routeList: list, net: QKDNetwork):
    # routeList: [{"sd": [1, 2, 4], "path": [1, 3, 2]}...]
    model = Model(name="provision_with_noise")
    T = model.continuous_var(name="T", lb=0)
    M = 1000
    e_dim1, e_dim2 = net.num_nodes + 1, net.num_nodes + 1

    e = model.binary_var_dict(
        [(u, v) for u in range(e_dim1) for v in range(e_dim2)], lb=0, name="p"
    )

    # 选一个发射方向
    for u in range(net.num_nodes):
        if net.G.nodes[u]["transmitter"]:
            sum_e = model.linear_expr()
            for v in net.G.neighbors(u):
                sum_e += e[u, v]
            model.add_constraint(sum_e == 1)

    def get_link_load(u, v):
        # 获取某个链路的负载
        count = 0
        for i in routeList:
            path = i["path"]
            l_path = len(path)
            for j in range(0, l_path - 1):
                link_s, link_e = path[j], path[j + 1]
                if [u, v] == [link_s, link_e] or [u, v] == [link_e, link_s]:
                    count += i["sd"][2] - i["sd"][3]
        count -= net.G[u][v]["remain_key_num"]
        return count

    # 压链路
    for u in range(net.num_nodes):
        if not net.G.nodes[u]["transmitter"]:
            continue
        for v in net.G.neighbors(u):
            sum_T = model.linear_expr()
            sum_T += get_link_load(u, v)
            sum_T -= e[u, v] * net.G.nodes[u]["transmitter_rate"]
            sum_T -= e[v, u] * net.G.nodes[v]["transmitter_rate"]
            model.add_constraint(sum_T <= T)

    model.minimize(T)
    sol = model.solve()

    ret_arr = []

    def has_load(u, v):
        # 判断某个链路是否有负载
        for i in routeList:
            path = i["path"]
            l_path = len(path)
            for j in range(0, l_path - 1):
                link_s, link_e = path[j], path[j + 1]
                if [u, v] == [link_s, link_e] or [u, v] == [link_e, link_s]:
                    return True
        return False

    if sol:
        # 打印值为1的e变量及其索引
        for u in range(e_dim1 - 1):
            for v in range(u + 1, e_dim2):
                if (u, v) in e and e[u, v].solution_value >= 0.99:
                    # if has_load(u, v):
                    ret_arr.append([u, v])
        # print("Variables e with value 1:")
        # print(ret_arr)
        return ret_arr

    return ret_arr
